import os
import wave
import numpy as np

def trim_wav(file_path, trim_start=0.05, trim_end=0.05):
    with wave.open(file_path, 'rb') as wav_file:
        params = wav_file.getparams()
        frames = wav_file.readframes(wav_file.getnframes())
    if trim_start == 0 and trim_end == 0:
        return frames, params
    sample_width = wav_file.getsampwidth()
    sample_rate = wav_file.getframerate()

    # 计算需要移除的采样数
    start_trim = int(trim_start * sample_rate)
    end_trim = int(trim_end * sample_rate)

    # 移除开头和结尾的采样数据
    trimmed_frames = frames[start_trim * sample_width: -end_trim * sample_width]

    return trimmed_frames, params

def concatenate_wav_files(folder_path, output_path, epoch, trim_start=0.05, trim_end=0.05):
    folder_path_epoch = os.path.join(folder_path, str(epoch))
    # 获取文件夹中所有的.wav文件
    wav_files = [f for f in os.listdir(folder_path_epoch) if f.endswith('.wav')]
    wav_files.sort()  # 按照文件名排序

    # 检查是否有.wav文件
    if not wav_files:
        print("文件夹中没有.wav文件")
        return

    # 读取第一个.wav文件的参数
    with wave.open(os.path.join(folder_path_epoch, wav_files[0]), 'rb') as first_wav:
        params = first_wav.getparams()
    output_file = 'epoch-' + str(epoch) + '.wav'
    output_file_path = os.path.join(output_path, output_file)
    # 创建一个新的.wav文件用于存储拼接后的音频
    with wave.open(output_file_path, 'wb') as output_wav:
        output_wav.setparams(params)

        # 逐个读取.wav文件并拼接
        for wav_file in wav_files:
            if not wav_file[:-4].isdigit():
                continue
            with wave.open(os.path.join(folder_path_epoch, wav_file), 'rb') as input_wav:
                # 读取音频数据
                trimmed_frames, _ = trim_wav(os.path.join(folder_path_epoch, wav_file), trim_start, trim_end)

                # frames = input_wav.readframes(input_wav.getnframes())
                # 写入拼接后的音频数据
                output_wav.writeframes(trimmed_frames)
    print('epoch-' + str(epoch) + '.wav'+"拼接完成")


if __name__ == '__main__':
    # 使用示例
    folder_path = "output/train_1s"  # 修改为你的文件夹路径
    output_path = "output/train_1s/audio"
    # for epoch in range(10, 1000, 10):
        # concatenate_wav_files(folder_path, output_path, epoch=epoch)
    concatenate_wav_files(folder_path, output_path, epoch=230, trim_start=0, trim_end=0)
